import { useMemo } from 'react';

type SetStateAction<S> = S | ((prevState: S) => S);
export function useCache<S = undefined>(): (action?: SetStateAction<S | undefined>) => S | undefined;
export function useCache<S>(initialValue?: S | (() => S)): (action?: SetStateAction<S>) => S;
export function useCache<S = undefined>(initialValue?: S | (() => S)) {
  const handler = useMemo(() => {
    let value: S | undefined = initialValue instanceof Function ? initialValue() : initialValue;
    return (...actions: [SetStateAction<S>]) => {
      const [action] = actions;
      if (actions.length <= 0) {
        return value;
      }

      if (action instanceof Function) {
        return (value = action(value as S));
      }

      return (value = action);
    };
  }, []);

  return handler;
}